# -*- coding: utf-8 -*-
'''
File: configuration_llama4.py
Time: 2025/9/15 14:50
IDE: PyCharm
Author: Barranzi An
email: awc19930818@outlook.com
github: https://github.com/La0bALanG
Barranzi's Blog: 私聊巴郎子要文档链接
requirement: Llama4模型配置文件，定义了视觉、文本和整体模型的配置参数
'''

# coding=utf-8
# Copyright 2025 The LLAMA4 and HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from ...configuration_utils import PretrainedConfig, layer_type_validation
from ...utils import logging

logger = logging.get_logger(__name__)


class Llama4VisionConfig(PretrainedConfig):
    r"""
    Llama4视觉模型的配置类，用于存储 [`Llama4VisionModel`] 的配置信息。
    根据指定的参数实例化Llama4视觉模型，定义模型架构。
    使用默认参数实例化配置将产生类似于Llama4 109B的配置。

    例如：[meta-llama/Llama-4-Scout-17B-16E](https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E)

    配置对象继承自 [`PretrainedConfig`] 并可用于控制模型输出。
    更多信息请阅读 [`PretrainedConfig`] 的文档。

    Args:
        hidden_size (`int`, *optional*, defaults to 768):
            编码器层和池化层的维度大小。这是视觉transformer的主要隐藏维度。
        hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
            编码器和池化器中的非线性激活函数（函数或字符串）。
            如果是字符串，支持 `"gelu"`, `"relu"`, `"selu"`, `"gelu_new"`, `"quick_gelu"`。
        num_hidden_layers (`int`, *optional*, defaults to 34):
            Transformer编码器中的隐藏层数量。决定模型的深度。
        num_attention_heads (`int`, *optional*, defaults to 16):
            Transformer编码器中每个注意力层的注意力头数量。
        num_channels (`int`, *optional*, defaults to 3):
            输入图像的通道数（RGB图像为3）。
        intermediate_size (`int`, *optional*, defaults to 5632):
            Transformer编码器中"中间"（通常称为前馈）层的维度。
            通常是hidden_size的4倍左右。
        vision_output_dim (`int`, *optional*, defaults to 7680):
            视觉模型输出的维度。包括带有中间层的transformer编码器
            和全局transformer编码器的输出。
        image_size (`int`, *optional*, defaults to 448):
            每个图像*块*的大小（分辨率）。输入图像会被resize到这个尺寸。
        patch_size (`int`, *optional*, defaults to 14):
            每个补丁(patch)的大小（分辨率）。图像会被分割成patch_size x patch_size的小块。
        norm_eps (`float`, *optional*, defaults to 1e-05):
            层归一化层使用的epsilon值，用于数值稳定性。
        vision_feature_layer (``, *optional*, defaults to -1):
            视觉特征提取的层索引，-1表示使用最后一层。
        vision_feature_select_strategy (`int`, *optional*, defaults to `"default"`):
            视觉特征选择策略，可以是"default"或"full"。
        initializer_range (`float`, *optional*, defaults to 0.02):
            用于初始化所有权重矩阵的截断正态分布初始化器的标准差。
        pixel_shuffle_ratio (`int`, *optional*, defaults to 0.5):
            像素混洗比率，用于视觉适配器中的上采样操作。
        projector_input_dim (`int`, *optional*, defaults to 4096):
            多模态投影器的输入维度。
        projector_output_dim (`int`, *optional*, defaults to 4096):
            多模态投影器的输出维度，通常与文本模型的隐藏维度一致。
        multi_modal_projector_bias (`bool`, *optional*, defaults to `False`):
            多模态投影器是否使用偏置项。
        projector_dropout (`float`, *optional*, defaults to 0.0):
            投影器中的dropout率，用于防止过拟合。
        attention_dropout (`float`, *optional*, defaults to 0.0):
            注意力层中的dropout率。
        rope_theta (`int`, *optional*, defaults to 10000):
            旋转位置编码(RoPE)的基础频率参数。
    """

    # 基础模型的张量并行计划，定义了各层的并行策略
    base_model_tp_plan = {
        # 注意力层的查询、键、值投影使用列式并行
        "model.layers.*.self_attn.q_proj": "colwise",
        "model.layers.*.self_attn.k_proj": "colwise",
        "model.layers.*.self_attn.v_proj": "colwise",
        # 输出投影使用行式并行
        "model.layers.*.self_attn.o_proj": "rowwise",
        # 视觉适配器的MLP层并行策略
        "vision_adapter.mlp.fc1": "colwise",
        "vision_adapter.mlp.fc2": "rowwise",
        # patch嵌入层使用列式复制并行
        "patch_embedding.linear": "colwise_rep",
    }

    # 模型类型标识
    model_type = "llama4_vision_model"
    # 基础配置键，用于嵌套配置
    base_config_key = "vision_config"

    def __init__(
            self,
            hidden_size: int = 768,  # 隐藏层维度
            hidden_act: str = "gelu",  # 激活函数
            num_hidden_layers: int = 34,  # 隐藏层数量
            num_attention_heads: int = 16,  # 注意力头数量
            num_channels: int = 3,  # 图像通道数
            intermediate_size: int = 5632,  # 中间层维度
            vision_output_dim: int = 7680,  # 视觉输出维度
            image_size: int = 448,  # 图像尺寸
            patch_size: int = 14,  # 补丁尺寸
            norm_eps: float = 1e-5,  # 归一化epsilon
            vision_feature_layer=-1,  # 视觉特征层索引
            vision_feature_select_strategy="default",  # 特征选择策略
            initializer_range: float = 0.02,  # 初始化范围
            pixel_shuffle_ratio=0.5,  # 像素混洗比率
            projector_input_dim=4096,  # 投影器输入维度
            projector_output_dim=4096,  # 投影器输出维度
            multi_modal_projector_bias=False,  # 投影器偏置
            projector_dropout=0.0,  # 投影器dropout
            attention_dropout=0.0,  # 注意力dropout
            rope_theta=10000,  # RoPE参数
            **kwargs,
    ):
        """初始化Llama4视觉配置"""
        # 保存所有配置参数
        self.hidden_size = hidden_size
        self.hidden_act = hidden_act
        self.num_hidden_layers = num_hidden_layers
        self.num_channels = num_channels
        self.intermediate_size = intermediate_size
        self.image_size = image_size
        self.vision_output_dim = vision_output_dim
        self.patch_size = patch_size
        self.norm_eps = norm_eps
        self.num_attention_heads = num_attention_heads
        self.initializer_range = initializer_range
        self.pixel_shuffle_ratio = pixel_shuffle_ratio
        self.projector_input_dim = projector_input_dim
        self.projector_output_dim = projector_output_dim
        self.multi_modal_projector_bias = multi_modal_projector_bias
        self.projector_dropout = projector_dropout
        self.attention_dropout = attention_dropout
        self.vision_feature_layer = vision_feature_layer
        self.vision_feature_select_strategy = vision_feature_select_strategy
        self.rope_theta = rope_theta

        # 调用父类初始化
        super().__init__(**kwargs)


class Llama4TextConfig(PretrainedConfig):
    r"""
    Llama4文本模型的配置类，用于存储 [`Llama4TextModel`] 的配置信息。
    根据指定的参数实例化Llama4文本模型，定义模型架构。
    使用默认参数实例化配置将产生类似于Llama4 109B的配置。

    例如：[meta-llama/Llama-4-Scout-17B-16E](https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E)

    配置对象继承自 [`PretrainedConfig`] 并可用于控制模型输出。
    更多信息请阅读 [`PretrainedConfig`] 的文档。

    Args:
        vocab_size (`int`, *optional*, defaults to 202048):
            Llama4文本模型的词汇表大小。定义了调用 [`Llama4TextModel`] 时
            `inputs_ids` 可以表示的不同token的最大数量。
        hidden_size (`int`, *optional*, defaults to 5120):
            嵌入和隐藏状态的维度。这是模型的主要表示维度。
        intermediate_size (`int`, *optional*, defaults to 8192):
            Transformer编码器中"中间"（通常称为前馈）层的维度。
        intermediate_size_mlp (`int`, *optional*, defaults to 16384):
            MLP层的中间维度，用于专家混合(MoE)架构。
        num_hidden_layers (`int`, *optional*, defaults to 48):
            Transformer编码器中的隐藏层数量。决定模型深度。
        num_attention_heads (`int`, *optional*, defaults to 40):
            Transformer编码器中每个注意力层的注意力头数量。
        num_key_value_heads (`int`, *optional*, defaults to 8):
            用于实现分组查询注意力(Grouped Query Attention)的键值头数量。
            如果未指定，将默认为 `num_attention_heads`。
        head_dim (`int`, *optional*, defaults to 128):
            每个注意力头的维度。通常是hidden_size除以num_attention_heads。
        hidden_act (`str` or `Callable`, *optional*, defaults to `"silu"`):
            编码器和池化器中的非线性激活函数（函数或字符串）。
        max_position_embeddings (`int`, *optional*, defaults to 131072):
            此模型可能使用的最大序列长度。定义了位置编码的范围。
        initializer_range (`float`, *optional*, defaults to 0.02):
            用于初始化所有权重矩阵的截断正态分布初始化器的标准差。
        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
            RMS归一化层使用的epsilon值，用于数值稳定性。
        use_cache (`bool`, *optional*, defaults to `True`):
            模型是否应该返回最后的键/值注意力。用于优化推理速度。
        pad_token_id (`int`, *optional*, defaults to 128004):
            填充token的ID。用于批处理时对齐序列长度。
        bos_token_id (`int`, *optional*, defaults to 1):
            句子开始token的ID。
        eos_token_id (`int`, *optional*, defaults to 2):
            句子结束token的ID。
        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
            是否绑定词嵌入权重。如果为True，输入和输出嵌入共享权重。
        rope_theta (`float`, *optional*, defaults to `500000.0`):
            RoPE嵌入的基础周期。更大的值支持更长的序列。
        attention_dropout (`float`, *optional*, defaults to 0.0):
            注意力层的dropout率。
        num_experts_per_tok (`int`, *optional*, defaults to 1):
            每个token激活的专家数量（用于MoE）。
        num_local_experts (`int`, *optional*, defaults to 16):
            总的本地专家数量（用于MoE）。
        moe_layers (`list[int]`, *optional*):
            使用MoE的层索引列表。如果为None，将自动生成。
        interleave_moe_layer_step (`int`, *optional*, defaults to 1):
            MoE层的交替步长。
        use_qk_norm (`bool`, *optional*, defaults to `True`):
            是否对查询和键使用归一化。
        output_router_logits (`bool`, *optional*, defaults to `False`):
            是否输出路由器logits（用于MoE）。
        router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
            路由器辅助损失系数（用于MoE负载均衡）。
        router_jitter_noise (`float`, *optional*, defaults to 0.0):
            路由器抖动噪声（用于MoE训练稳定性）。
        rope_scaling (`Dict`, *optional*):
            RoPE嵌入的缩放配置字典。注意：如果应用新的rope类型
            并希望模型在更长的 `max_position_embeddings` 上工作，
            建议相应地更新此值。
            预期内容：
                `rope_type` (`str`):
                    要使用的RoPE子变体。可以是['default', 'linear', 'dynamic', 'yarn', 'longrope',
                    'llama3']之一，'default'是原始RoPE实现。
                `factor` (`float`, *optional*):
                    除'default'外的所有rope类型都使用。应用于RoPE嵌入的缩放因子。
                    在大多数缩放类型中，因子x将使模型能够处理长度为x *
                    原始最大预训练长度的序列。
                `original_max_position_embeddings` (`int`, *optional*):
                    用于'dynamic', 'longrope'和'llama3'。预训练期间使用的原始最大位置嵌入。
                `attention_factor` (`float`, *optional*):
                    用于'yarn'和'longrope'。应用于注意力计算的缩放因子。
                `beta_fast` (`float`, *optional*):
                    仅用于'yarn'。设置线性斜坡函数中外推边界的参数。
                `beta_slow` (`float`, *optional*):
                    仅用于'yarn'。设置线性斜坡函数中插值边界的参数。
                `short_factor` (`list[float]`, *optional*):
                    仅用于'longrope'。应用于短上下文的缩放因子。
                `long_factor` (`list[float]`, *optional*):
                    仅用于'longrope'。应用于长上下文的缩放因子。
                `low_freq_factor` (`float`, *optional*):
                    仅用于'llama3'。应用于RoPE低频分量的缩放因子。
                `high_freq_factor` (`float`, *optional*):
                    仅用于'llama3'。应用于RoPE高频分量的缩放因子。
        no_rope_layers (`list[int]`, *optional*):
            至少与模型中的层数相同长度的列表。
            索引位置的 `1` 表示相应层将使用RoPE，
            而 `0` 表示它是NoPE层。
        no_rope_layer_interval (`int`, *optional*, defaults to 4):
            如果 `no_rope_layers` 为 `None`，将使用每隔
            `no_rope_layer_interval` 层创建一个NoPE层。
        attention_chunk_size (`int`, *optional*, defaults to 8192):
            注意力分块大小，用于长序列的内存优化。
        layer_types (`list`, *optional*):
            每层的注意力模式。可以是"full_attention"或"chunked_attention"。
        attn_temperature_tuning (`bool`, *optional*, defaults to `True`):
            是否根据序列长度为每个查询token动态缩放注意力温度。
            推荐用于长序列（例如，>32k tokens）以保持稳定的输出结果。
        floor_scale (`int`, *optional*, defaults to 8192):
            注意力温度调节的最小缩放值。
        attn_scale (`float`, *optional*, defaults to 0.1):
            注意力缩放因子。
    """

    # 模型类型标识
    model_type = "llama4_text"
    # 推理时忽略的键
    keys_to_ignore_at_inference = ["past_key_values"]

    # 基础模型张量并行计划
    base_model_tp_plan = {
        # 注意力层的投影矩阵并行策略
        "layers.*.self_attn.q_proj": "colwise",  # 查询投影：列式并行
        "layers.*.self_attn.k_proj": "colwise",  # 键投影：列式并行
        "layers.*.self_attn.v_proj": "colwise",  # 值投影：列式并行
        "layers.*.self_attn.o_proj": "rowwise",  # 输出投影：行式并行

        # 层归一化权重：序列并行
        "layers.*.input_layernorm.weight": "sequence_parallel",
        "layers.*.post_attention_layernorm.weight": "sequence_parallel",
        "norm.weight": "sequence_parallel",

        # 前馈网络的并行策略
        # 共享专家层
        "layers.*.feed_forward.shared_expert.gate_proj": "local_colwise",
        "layers.*.feed_forward.shared_expert.up_proj": "local_colwise",
        "layers.*.feed_forward.shared_expert.down_proj": "local_rowwise",

        # MoE专家层
        "layers.*.feed_forward.experts.gate_up_proj": "local_packed_rowwise",
        "layers.*.feed_forward.experts.down_proj": "local_colwise",
        "layers.*.feed_forward.experts": "local",

        # 标准前馈层
        "layers.*.feed_forward.gate_proj": "local_colwise",
        "layers.*.feed_forward.up_proj": "local_colwise",
        "layers.*.feed_forward.down_proj": "local_rowwise",
        "layers.*.feed_forward": "gather",
    }

    # 基础模型专家并行计划
    base_model_ep_plan = {
        # 注意力层并行策略
        "layers.*.self_attn.q_proj": "colwise",
        "layers.*.self_attn.k_proj": "colwise",
        "layers.*.self_attn.v_proj": "colwise",
        "layers.*.self_attn.o_proj": "rowwise",

        # MoE专家层的专家并行策略
        "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm",
        "layers.*.feed_forward.experts.down_proj": "grouped_gemm",
        "layers.*.feed_forward.experts": "gather",

        # 标准前馈层
        "layers.*.feed_forward.gate_proj": "local_colwise",
        "layers.*.feed_forward.up_proj": "local_colwise",
        "layers.*.feed_forward.down_proj": "local_rowwise",
        "layers.*.feed_forward.router": "ep_router",
    }

    def __init__(
            self,
            vocab_size=202048,  # 词汇表大小
            hidden_size=5120,  # 隐藏层维度
            intermediate_size=8192,  # 中间层维度
            intermediate_size_mlp=16384,  # MLP中间层维度
            num_hidden_layers=48,  # 隐藏层数量
            num_attention_heads=40,  # 注意力头数量
            num_key_value_heads=8,  # 键值头数量（GQA）
            head_dim=128,  # 注意力头维度
            hidden_act="silu",  # 隐藏层激活函数
            max_position_embeddings=4096 * 32,  # 最大位置嵌入（128K）
            initializer_range=0.02,  # 初始化范围
            rms_norm_eps=1e-5,  # RMS归一化epsilon
            use_cache=True,  # 是否使用缓存
            pad_token_id=None,  # 填充token ID
            bos_token_id=1,  # 开始token ID
            eos_token_id=2,  # 结束token ID
            tie_word_embeddings=False,  # 是否绑定词嵌入
            rope_theta=500000,  # RoPE基础频率
            attention_dropout=0.0,  # 注意力dropout
            num_experts_per_tok=1,  # 每token专家数
            num_local_experts=16,  # 本地专家总数
            moe_layers=None,  # MoE层列表
            interleave_moe_layer_step=1,  # MoE层间隔
            use_qk_norm=True,  # 是否使用QK归一化
            output_router_logits=False,  # 是否输出路由器logits
            router_aux_loss_coef=0.001,  # 路由器辅助损失系数
            router_jitter_noise=0.0,  # 路由器抖动噪声
            rope_scaling=None,  # RoPE缩放配置
            no_rope_layers=None,  # 不使用RoPE的层
            no_rope_layer_interval=4,  # NoPE层间隔
            attention_chunk_size=8192,  # 注意力分块大小
            layer_types=None,  # 层类型列表
            attn_temperature_tuning=True,  # 注意力温度调节
            floor_scale=8192,  # 最小缩放值
            attn_scale=0.1,  # 注意力缩放因子
            **kwargs,
    ):
        """初始化Llama4文本配置"""
        # 调用父类初始化，设置基本token参数
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )

        # 保存注意力相关配置
        self.attn_temperature_tuning = attn_temperature_tuning
        self.attn_scale = attn_scale
        self.floor_scale = floor_scale

        # 保存基本模型配置
        self.vocab_size = vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.intermediate_size_mlp = intermediate_size_mlp
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.rope_scaling = rope_scaling
        self.attention_bias = False  # Llama不使用注意力偏置

        # 向后兼容性处理
        if num_key_value_heads is None:
            num_key_value_heads = num_attention_heads

        self.num_key_value_heads = num_key_value_heads
        self.hidden_act = hidden_act
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps
        self.use_cache = use_cache
        self.rope_theta = rope_theta
        self.attention_dropout = attention_dropout
        # 计算注意力头维度，如果未指定则自动计算
        self.head_dim = head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads
        self.use_qk_norm = use_qk_norm

        # MoE相关配置
        self.num_experts_per_tok = num_experts_per_tok
        self.num_local_experts = num_local_experts
        self.output_router_logits = output_router_logits
        self.router_aux_loss_coef = router_aux_loss_coef
        self.router_jitter_noise = router_jitter_noise

        # NoPE层配置的向后兼容性
        if no_rope_layers == []:
            no_rope_layers = None

        # 创建默认的NoPE层配置：每隔no_rope_layer_interval设置一个NoPE层
        default_no_rope_layers = [
            int((layer_idx + 1) % no_rope_layer_interval != 0)
            for layer_idx in range(self.num_hidden_layers)
        ]

        self.no_rope_layers = no_rope_layers if no_rope_layers else default_no_rope_layers

        # MoE层配置
        self.interleave_moe_layer_step = interleave_moe_layer_step
        # 如果未指定MoE层，则按步长自动生成
        self.moe_layers = (
            moe_layers
            if moe_layers is not None
            else list(range(interleave_moe_layer_step - 1, num_hidden_layers, interleave_moe_layer_step))
        )
        self.attention_chunk_size = attention_chunk_size

        # 层类型配置：根据是否使用RoPE决定注意力类型
        self.layer_types = layer_types
        if layer_types is None:
            self.layer_types = [
                "chunked_attention" if no_rope else "full_attention"
                for no_rope in self.no_rope_layers
            ]
        # 验证层类型配置的有效性
        layer_type_validation(self.layer_types)


class Llama4Config(PretrainedConfig):
    r"""
    Llama4模型的主配置类，用于存储 [`Llama4Model`] 的配置信息。
    根据指定的参数实例化Llama4模型，定义模型架构。
    使用默认参数实例化配置将产生类似于Llama4 109B的配置。

    例如：[meta-llama/Llama-4-Scout-17B-16E](https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E)

    配置对象继承自 [`PretrainedConfig`] 并可用于控制模型输出。
    更多信息请阅读 [`PretrainedConfig`] 的文档。

    这是一个多模态配置类，整合了视觉和文本两个子模型的配置，
    并添加了多模态交互所需的特殊token配置。

    Args:
        vision_config (`Llama4VisionConfig`, *optional*):
            Llama4视觉模块的配置对象。如果为None，将使用默认的视觉配置。
            控制视觉编码器的架构参数，如层数、注意力头数、图像尺寸等。
        text_config (`Llama4TextConfig`, *optional*):
            Llama4文本模块的配置对象。如果为None，将使用默认的文本配置。
            控制语言模型的架构参数，如词汇表大小、隐藏维度、层数等。
        boi_token_index (`int`, *optional*, defaults to 200080):
            图像开始token的索引，用于包装图像提示。
            在输入序列中标记图像内容的开始位置，帮助模型理解多模态输入的结构。
        eoi_token_index (`int`, *optional*, defaults to 200081):
            图像结束token的索引，用于包装图像提示。
            在输入序列中标记图像内容的结束位置，与boi_token配对使用。
        image_token_index (`int`, *optional*, defaults to 200092):
            图像token的索引，用于编码图像提示。
            这是实际的图像占位符token，在文本序列中代表图像特征的位置。
            在前向传播时，这些token会被实际的图像特征替换。
        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
            模型的输入和输出词嵌入是否应该绑定（共享权重）。
            如果为True，输入嵌入层和输出投影层将共享相同的权重矩阵，
            这可以减少参数量并可能提高训练稳定性。

    使用示例:
    ```python
    >>> from transformers import Llama4Model, Llama4Config

    >>> # 初始化一个Llama4 7B风格的配置
    >>> configuration = Llama4Config()

    >>> # 从Llama4 7B风格配置初始化模型
    >>> model = Llama4Model(configuration)

    >>> # 访问模型配置
    >>> configuration = model.config
    ```
    """

    # === 类属性定义 ===

    # 模型类型标识符，用于HuggingFace模型注册和识别
    model_type = "llama4"

    # 属性映射表：将不同的命名约定映射到统一的内部属性名
    # 这样可以兼容不同的命名方式，提高API的灵活性
    attribute_map = {
        "image_token_id": "image_token_index",  # 图像token ID的别名
        "boi_token_id": "boi_token_index",  # 图像开始token ID的别名
        "eoi_token_id": "eoi_token_index",  # 图像结束token ID的别名
    }

    # 子配置映射：定义了嵌套配置的类型
    # 当从字典创建配置时，这些字段会被自动转换为相应的配置对象
    sub_configs = {
        "text_config": Llama4TextConfig,  # 文本配置类
        "vision_config": Llama4VisionConfig  # 视觉配置类
    }

    # 基础模型的张量并行计划
    # 定义了多模态投影器的并行策略
    base_model_tp_plan = {
        "multi_modal_projector.linear_1": "colwise_rep",  # 投影器第一层使用列式复制并行
    }

    def __init__(
            self,
            vision_config=None,  # 视觉配置
            text_config=None,  # 文本配置
            boi_token_index=200080,  # 图像开始token索引
            eoi_token_index=200081,  # 图像结束token索引
            image_token_index=200092,  # 图像token索引
            tie_word_embeddings=False,  # 是否绑定词嵌入
            **kwargs,  # 其他配置参数
    ):
        """
        初始化Llama4配置对象

        该方法负责：
        1. 处理和验证视觉配置
        2. 处理和验证文本配置
        3. 设置多模态特殊token
        4. 调用父类初始化
        """

        # === 第一步：处理视觉配置 ===
        # 支持三种输入方式：None、字典、配置对象
        if vision_config is None:
            # 如果未提供视觉配置，使用默认配置
            self.vision_config = Llama4VisionConfig()
            logger.info("vision_config is None, using default llama4 vision config")
        elif isinstance(vision_config, dict):
            # 如果提供的是字典，将其转换为配置对象
            self.vision_config = Llama4VisionConfig(**vision_config)
        elif isinstance(vision_config, Llama4VisionConfig):
            # 如果已经是配置对象，直接使用
            self.vision_config = vision_config

        # === 第二步：设置多模态特殊token ===
        # 这些token用于在文本序列中标记和表示图像内容
        self.boi_token_index = boi_token_index  # 图像开始标记
        self.eoi_token_index = eoi_token_index  # 图像结束标记
        self.image_token_index = image_token_index  # 图像占位符标记

        # === 第三步：处理文本配置 ===
        # 同样支持三种输入方式：None、字典、配置对象
        if text_config is None:
            # 如果未提供文本配置，使用默认配置
            self.text_config = Llama4TextConfig()
            logger.info("text_config is None, using default llama4 text config")
        elif isinstance(text_config, dict):
            # 如果提供的是字典，将其转换为配置对象
            self.text_config = Llama4TextConfig(**text_config)
        elif isinstance(text_config, Llama4TextConfig):
            # 如果已经是配置对象，直接使用
            self.text_config = text_config

        # === 第四步：调用父类初始化 ===
        # 传递词嵌入绑定参数和其他配置参数
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)


# === 模块导出列表 ===
# 定义了该模块对外暴露的公共API
# 这样可以控制import *时导入的内容，保持API的清洁
__all__ = [
    "Llama4Config",  # 主配置类：整合视觉和文本配置的多模态配置
    "Llama4TextConfig",  # 文本配置类：专门用于语言模型部分的配置
    "Llama4VisionConfig"  # 视觉配置类：专门用于视觉编码器部分的配置
]
